﻿using FlyIL;
using FlyIL.Advance;
using FlyIL.Lifter;
using FlyRemote.Base;
using FlyRemote.Remote;
using FlyRemote.Scanning.Searcher;
using FuckVmp.MemoryPE;

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;
using System.Threading.Tasks;

namespace FuckVmp.VMP
{
    public class Detector
    {

        [DllImport("msvcrt.dll", CallingConvention = CallingConvention.Cdecl)]
        public static extern IntPtr memchr(IntPtr ptr, byte chr, int size);





        string[] vSections = new string[] { "vmp" };

        ReadGroup Group { get; set; }

        public IReader Reader { get; set; }

        public IReport Report { get; set; }

        public ISymbolHandler Symbol { get; set; }


        PEInfo Info { get; set; }

        public HashSet<VMIat> VMIats { get; } = new HashSet<VMIat>();

        public HashSet<VMEntry> VMEntries { get; } = new HashSet<VMEntry>();

        public HashSet<VMUnknow> VMUnknows { get; } = new HashSet<VMUnknow>();


        public Detector(IntPtr PHandle, ReadGroup group)
        {

            


            Group = group;

            if (Group.CurrentValues == null)
                Group.ReadCurrentValues();

            this.Info = PEInfo.Load(Group.BaseAddress, Group.CurrentValues);


        }


        /// <summary>
        /// 寻找 跨 section 的跳转 或者 call 
        /// </summary>
        IList<SnapshotCell> SearchCore(ReadGroup group, int start, int end, byte firstByte)
        {
            List<SnapshotCell> list = new List<SnapshotCell>(0x10);

            long groupS = (long)Marshal.UnsafeAddrOfPinnedArrayElement(group.CurrentValues, 0);
            long groupE = (long)Marshal.UnsafeAddrOfPinnedArrayElement(group.CurrentValues, group.RegionSize);

            long ptr = (long)Marshal.UnsafeAddrOfPinnedArrayElement(group.CurrentValues, start);
            long ptrEnd = ((long)ptr + end);
            long cur = ptr;

            for (; cur < ptrEnd; cur++)
            {
                cur = (long)memchr((IntPtr)cur, firstByte, (int)(ptrEnd - cur));

                if (cur == 0)
                    break;

                int next = Marshal.ReadInt32((IntPtr)(cur + 1));
                long nextA = cur + 5 + next;

                // 如果是 本区间的跳转 ， 或者Call 
                if (ptr <= nextA && nextA <= ptrEnd)
                    continue;

                // 如果超出了本模块  
                if (nextA < groupS || groupE < nextA)
                    continue;

                PEInfo.IMAGE_SECTION_HEADER section = null;

                for (int i = 0; i < Info.Sections.Length; i++)
                {
                    var u = Info.Sections[i];
                    uint secS = i == 0 ? 0 : Info.Sections[i - 1].VirtualAddress;
                    if (secS < (nextA - groupS) && (nextA - groupS) < (secS + u.PhysicalAddress))
                    {
                        section = u;
                        break;
                    }
                }

                /// 跳转的目标区段 是 vmp0 或者 vmp1 
                if (vSections.Any(o => section != null && section.name.Contains(o)))
                {
                    var c = new SnapshotCell(group, (int)(cur - groupS), 5);
                    list.Add(c);
                }

            }

            return list;
        }

        /// <summary>
        /// 寻找跨区跳转 ， 即vmp 中的 VEntry 或 VIat 
        /// </summary>
        /// <returns></returns>
        List<SnapshotCell> SearchSection(string sectionName = ".text")
        {

            /*
             *   从 .text 区段 -> vmp区段 ， 使用的是 call 指令 
             */

            IntPtr groupStart = Marshal.UnsafeAddrOfPinnedArrayElement(Group.CurrentValues, 0);
            IList<SnapshotCell> cells = BytesSearcher.Search(Group, 0xE9, (IntPtr a) => 5);
            List<SnapshotCell> list = new List<SnapshotCell>();

            for (int i = 0; i < Info.Sections.Length; i++)
            {

                var u = Info.Sections[i];
                if (u.SizeOfRawData > 0)
                    continue;

                // 可以执行的 
                if ((u.Characteristics & PEInfo.IMAGE_SCN_MEM_EXECUTE) == 0)
                    continue;

                if (u.name.Trim() != sectionName)
                    continue;

                /// 该区段的 内存起始位置  ,等于 前一个 区段的 VirtualAddress 
                int start = i == 0 ? 0 : (int)Info.Sections[i - 1].VirtualAddress;
                int end = (int)(u.VirtualAddress + u.PhysicalAddress);

                /*
                 * 目前可知 vmiat 均是伪装成 call xxxx  (跨区段跳转 
                 * 但是 vm入口 ， 可能会 伪装成 call jmp jcc 等多种形式， 目前只处理 call jmp 
                 * 
                 */

                list.AddRange(SearchCore(Group, start, end, 0xE8));
                list.AddRange(SearchCore(Group, start, end, 0xE9));

            }

            // 搜索到的 跨区段 跳转 , 有一部分是 对导入表的变形  

            return list;
        }

        public void DetectIL(ulong address)
        {


            try
            {
                Context vcon = null;

                var lifter = new Lifters.Intel(this.Info.Is64, Reader);
                vcon = new Context(this.Info.Is64, Group.BaseAddress)
                {
                    Reader = Reader
                };

                File.Delete(@"d:\\dump.asm");
                if (!lifter.Collect(vcon,
                   option: new Lifters.IntelOption
                   {
                       ///note: vmp 为每一处 导入表的调用, 均设计一段垃圾代码
                       /// 因此 viat的 识别, 只能一个个化简,无法取巧
                       Simplify = (Context ct) => ct.SimplifyGroup(),
                       CallAsjmp = true,
                       TrustJcc = true
                   },
                   entry: address, verifyAsm: true))
                {
                    return;
                }

                vcon.RebuildTags();
                vcon.ClearBranch();
                // vcon.DumpFlow();
                var inst = vcon.Last;

                /// 跳转到固定的地址 
                if (inst.Opt == Opt.jmp && inst.O1.IsImm)
                {

                    /// 如果是 32位 是入栈 4 , 64位 则是入栈 8
                    if (vcon.CountSpVariety >= (vcon.Is64 ? -8 : -4))
                    {

                        /// 跳出到模块外
                        if (Reader.IsOutOfModule(inst.O1.u64))
                        {
                            var next = vcon.Last.O1.u64;

                            var sym = Symbol.GetSymInfo(next);
                            var iat = new VMIat { Entry = vcon.EntryPoint, JmpAddr = next, Note = $"{sym?.Module}.{sym?.Name}" };

                            Report?.Comment(iat.Entry, iat.ToString());

                            lock (this)
                                VMIats.Add(iat);
                            return;
                        }

                    }

                }



                var entry = new VMEntry(Group, address, Reader, vcon);

                Report?.Comment(entry.Entry, entry.ToString());
                lock (this)
                    VMEntries.Add(entry);
                return;
            }
            catch (Exception e)
            {
                /// 如果不能识别 , 或者化简中 遇到了异常 
                lock (this)
                    VMUnknows.Add(new VMUnknow { Entry = address });

                Report?.Log($"unknow {address:x} {e}");
            }

        }


        public void Search()
        {

            var list = SearchSection();


            /// 这里使用并发搜索 
            Parallel.For(0, list.Count, i =>
            {

                DetectIL(list[i].BaseAddress);

            });

        }


    }
}
